module Zeno.Solver
(
  solver
)
where

import Prelude ()
import StdImports
import qualified Data.Set as Set

import Zeno.Core
import Zeno.Type
import Zeno.Term
import Zeno.Clause
import Zeno.Function
import Zeno.ProofLogic ( Proof (..) )
import qualified Zeno.ProofLogic as PL 
import Zeno.Debug

{- 
TODO

Test whether "contradiction" ever needs a depth greater than one
  but it makes sense that it would
  
Do things in phases?
  like after induction there is no generalisation?
  after splitting there is no induction?
  
Remember to add kterm hypotheses to the conditions
  for elem-min
  
Theory about removal of multiple linked induction hypotheses on application

Theory about only adding a hypothesis one the "first" inductive step of a chain
  that is to say do not add a hypothesis if hypotheses already exist
  
Try merging case-splitting back into induction
  just remember to apply function definitions whilst finding critical terms
  
Try iding each case-analysis and never use case-analysis to do induction on
  a child of a variable that was created with this analysis
  
Try making induction deterministic like case-splitting

Fix the randomness so every method returns Solver [ZProof]
  no more splitting the random generator though I guess
  could split it down each different step though...
  
Maybe remove all hypothesis after hgn-ing
  pretty sure this won't work

Don't need the antecedent removal system anymore it seems,
  woopdy doo
  
Try moving function definitions into the function variables themselves,
  might save a bit of time removing the lookup
  
*******
Express ALL/EXISTS in Haskell itself,
  i.e. a non-terminating function which builds up every single 
    different value of a type
  see if we can prove things about this,
    a generalised notion of induction
*******
-}

data SolverState
  = SolverState   { sZeno :: !Zeno,
                    sGoal :: !ZClause,
                    sHypotheses :: ![ZClause],
                    sDoGeneralisation :: !Bool,
                    sIndSteps :: !Int,
                    sConSteps :: !Int,
                    sRandom :: !StdGen }
    
type Solver = Reader SolverState 
type StatefulSolver = State SolverState
type SolverStep = Solver [ZProof]

liftZenoReader :: ZenoReader a -> Solver a
liftZenoReader zr = do
  zeno <- sZeno <$> ask
  return $ runReader zr zeno
  
liftZenoState :: ZenoState a -> StatefulSolver a
liftZenoState zs = do
  zeno <- sZeno <$> get
  let (a, zeno') = runState zs zeno
  modify $ \s -> s { sZeno = zeno' }
  return a
  
solver :: (Int, Int) -> String -> StdGen -> ZenoReader ZProof
solver (ind, con) lemma_name random = do
  func <- lookupFunctionName lemma_name
  lemmas <- clausifyFunction func
  proofs <- mapM solver' lemmas
  return $ case proofs of
    [] -> Unknown
    [proof] -> proof
    proofs -> PL.allProvenStep (PL.ProofStep "proofs" Nothing []) proofs
  where
    solver' :: ZClause -> ZenoReader ZProof
    solver' lemma = do
      zeno <- ask
      let state =
            SolverState { sZeno = zeno,
                          sGoal = lemma { clauseVariables = [] },
                          sHypotheses = mempty,
                          sDoGeneralisation = True,
                          sIndSteps = ind,
                          sConSteps = con,
                          sRandom = random }
      return $ runReader (addStep "goal" [] solve) state
    
solve :: Solver ZProof
solve = do
  Clause [] eq conds <- sGoal <$> ask
  let eq' = mapPair (substituteWithin (Set.fromList conds)) eq 
  (goal', changed) <- liftZenoReader $ 
    applyTermDefinitions conds (Clause [] eq' conds)
    
  let steps = concat <$> randomSequence
                [solve_hgn, solve_hyp, solve_fac, 
                 solve_gen, solve_cse, solve_ind]
      proof = checkGoal $ PL.anyProven <$> steps
  if not changed
    then proof
    else localGoal goal' $ addStep "def" [] $ proof

solve_hyp :: Solver [ZProof]
solve_hyp = enableGeneralisation $ cannotDisprove $ do
  applyHypotheses maybeApplyHyp
  where
    maybeApplyHyp :: ZClause -> Solver [ZProof]
    maybeApplyHyp hyp = do
      goal@(l_goal, r_goal) <- clauseEquality <$> sGoal <$> ask
      let apply_left = findApplications l_goal hyp
          apply_right = findApplications r_goal (flipClauseEquality hyp)
          l_eqs = map (\t -> (t, r_goal)) apply_left
          r_eqs = map (\t -> (l_goal, t)) apply_right
          all_eqs = l_eqs ++ r_eqs
      hyp_s <- liftZenoReader $ showZClause hyp
      let step = "hyp " ++ hyp_s
      mapM (\eq -> localGoalEquality eq $ addStep step [] solve) all_eqs
    
    findApplications :: ZTerm -> ZClause -> [ZTerm]
    findApplications goal_term 
                     (Clause vars (from_hyp, to_hyp) []) = 
      map applyUnifier $ 
        filter (completeUnifier vars) $ 
          allUnifiers from_hyp goal_term
      where 
        applyUnifier :: ZTermSubstitution -> ZTerm
        applyUnifier uni = 
          let hyp_uni = mapPair (substituteWithin uni) (from_hyp, to_hyp)
          in substituteWithin (Set.singleton hyp_uni) goal_term
          
          
solve_fac :: Solver [ZProof]
solve_fac = cannotDisprove $ do
  goal <- sGoal <$> ask
  let (l_term,  r_term) = clauseEquality goal
      (l_func : l_args) = flattenTerm l_term
      (r_func : r_args) = flattenTerm r_term
  if l_func /= r_func
    then return [Unknown]
    else do
      step <- proofstep "fct" []
      let factor_goals = l_args `zip` r_args
          factor_proofs = mapM 
            (\eq -> localGoalEquality eq $ addStep "fac" [] solve) factor_goals
      return <$> PL.allProvenStep step <$> factor_proofs
      
      
solve_hgn :: Solver [ZProof]
solve_hgn = applyHypotheses generaliseHyp
  where
    generaliseHyp :: ZClause -> Solver [ZProof]
    generaliseHyp (Clause vars eq@(l_hyp, r_hyp) []) 
      | not (constructorTerm r_hyp) = return [Unknown]
      | otherwise = do
          goal_term <- fst <$> clauseEquality <$> sGoal <$> ask 
          
          let possibleTerm :: ZTerm -> Bool
              possibleTerm term = 
                not (isVar term) &&
                functionTerm term &&
                destructibleTerm term &&
                length (variables term) == length vars + 1
    
          sequence $ do
            possible <- filter possibleTerm (strictlyWithinList l_hyp)
            uni <- allUnifiers possible goal_term
            if not (completeUnifier vars uni) 
              then return (return Unknown)
              else do
                let eq' = mapPair (substituteWithin uni) eq
                    genr = substituteWithin uni possible
                return $ addCond eq' $ 
                         addStep "hgn" [] $ 
                         applyGeneralisation genr
      

solve_gen :: Solver [ZProof]
solve_gen = generalisation $ do
  inds <- filter variableTerm <$> criticalTerms
  goal <- sGoal <$> ask
  let superterms = concatMap (duplicateSuperterms (withinList goal)) inds
  mapM applyGeneralisation superterms
  where  
    duplicateSuperterms :: [ZTerm] -> ZTerm -> [ZTerm]
    duplicateSuperterms all_terms sub_term = 
      removeFactors all_duplicates
      where
        all_duplicates :: [ZTerm]
        all_duplicates = filter superterm (duplicates all_terms)
        
        superterm :: ZTerm -> Bool
        superterm t = destructibleTerm t && isApp t 
          && t `contains` sub_term && functionTerm t
        
solve_cse :: Solver [ZProof]
solve_cse = disableGeneralisation $ takeIndStep $ cannotDisprove $ do
  splits <- filter isApp <$> criticalTerms
  mapM goSplit splits
  where
    goSplit :: ZTerm -> Solver ZProof
    goSplit term 
      | not (destructibleTerm term) = return Unknown
      | otherwise = do
          splits <- proveOverConstructors goSplit' (termType term)
          term_s <- liftZenoReader $ showZTerm term
          step <- proofstep "cse" []
          return $ PL.allProvenStep step splits
      where
        goSplit' :: (ZTerm, [ZVar]) -> Solver ZProof
        goSplit' (kterm, _) = do
          Clause [] eq conds <- sGoal <$> ask
          let goal = Clause [] eq ((term, kterm) : conds)
          
          term_s <- liftZenoReader $ showZTerm term
          kterm_s <- liftZenoReader $ showZTerm kterm
          let step = "cse " ++ term_s ++ " => " ++ kterm_s
          localGoal goal $ addStep step [] solve


solve_ind :: Solver [ZProof]
solve_ind = disableGeneralisation $ takeIndStep $ do
  inds <- map fromVar . filter variableTerm <$> criticalTerms
  mapM goInductive inds
  where
    goInductive :: ZVar -> Solver ZProof
    goInductive var 
      | not (destructibleVar var) = return Unknown
      | otherwise = do
          inds <- proveOverConstructors goInductive' (zVarType var)
          var_s <- liftZenoReader $ showZVar var
          step <- proofstep "ind" []
          return $ PL.allProvenStep step inds
      where
        goInductive' :: (ZTerm, [ZVar]) -> Solver ZProof
        goInductive' (ind, ind_vars) = do
          goal <- sGoal <$> ask
          let ind_sub = Set.singleton (Var var, ind)
          
              neededCondition :: ZEquality -> Bool
              neededCondition (t1, t2) = 
                var `elem` toList t1 || var `elem` toList t2
              
              hyp_goal = goal 
                { clauseConditions = 
                    filter neededCondition (clauseConditions goal) }
              
              createHyp :: ZVar -> StatefulSolver ZClause
              createHyp ind_var = liftZenoState $ do
                let others = filter (/= var) (variables hyp_goal)
                others' <- mapM freshVariable others
                let sublist = (var, ind_var) : (others `zip` others')
                    sub = Set.fromList $ map (Var *** Var) sublist
                    hyp = substituteWithin sub hyp_goal
                return $ hyp { clauseVariables = others' }
                
          var_s <- liftZenoReader $ showZVar var
          ind_s <- liftZenoReader $ showZTerm ind
          let step = "ind " ++ var_s ++ " => " ++ ind_s
                
          stateful (mapM createHyp ind_vars) $ \hyps -> do
            substituteTerms ind_sub $ 
              addHypotheses hyps $
                addStep step hyps solve

proveOverConstructors :: ((ZTerm, [ZVar]) -> Solver a) -> ZType -> Solver [a] 
proveOverConstructors prove ztype = do
  let ((TypeVar dtype@(ZDataType _ dt_args cons)) : ty_args) = flattenAppType ztype
      bindings = Set.fromList $ (map TypeVar dt_args) `zip` ty_args  
  stateful (mapM (liftZenoState . instantiate bindings) cons) (mapM prove)

applyGeneralisation :: ZTerm -> Solver ZProof
applyGeneralisation term = do
  stateful (liftZenoState newIdS) $ \new_id -> do
    let new_var = ZVar new_id (termType term) ZVariable
        term_sub = Set.singleton (term, Var new_var)
    term_s <- liftZenoReader $ showZTerm term
    var_s <- liftZenoReader $ showZVar new_var
    let step = "gen " ++ term_s ++ " => " ++ var_s
    substituteTerms term_sub $
      addStep step [] solve
              
completeUnifier :: [ZVar] -> ZTermSubstitution -> Bool
completeUnifier vars uni =
  Set.fromList (map Var vars) == Set.mapMonotonic fst uni
              
criticalTerms :: Solver [ZTerm]
criticalTerms = do
  Clause [] (l_term, r_term) conds <- sGoal <$> ask
  let all_terms = l_term : r_term : (map fst conds)
  filter destructibleTerm . nub <$> mapMaybeM criticalTerm all_terms
  where
    criticalTerm :: ZTerm -> Solver (Maybe ZTerm)
    criticalTerm term = do
      maybe_expr <- liftZenoReader $ applyFunction [] term
      case maybe_expr of
        Nothing -> 
          return (Just term)
        Just (Value {}) -> 
          return Nothing
        Just (CaseOf term' _) -> do
          if term `contains` term'
            then criticalTerm term'
            else return (Just term')

-- Applicable hypotheses are those with all conditions fulfilled
--   the given hypotheses should be conditionless
applyHypotheses :: (ZClause -> Solver [ZProof]) -> Solver [ZProof]
applyHypotheses apply = do
  hyps <- sHypotheses <$> ask
  concat <$> concatMapM applyHyp hyps
  where
    applyHyp :: ZClause -> Solver [[ZProof]]
    applyHyp cls = do
      apps <- applicable cls
      removeHypothesis cls (mapM apply apps)
  
    applicable :: ZClause -> Solver [ZClause]
    applicable cls@(Clause _ _ []) = return [cls]
    applicable (Clause c_vars c_eq c_conds) = do
      unis <- mapM removeCondition c_conds
      let unis' = map Set.unions (sequence unis)
      return $ do
        uni <- unis'
        let vars = map (fromVar . fst) (toList uni)
            c_vars' = filter (not . flip elem vars) c_vars
            c_eq' = mapPair (substituteWithin uni) c_eq
        return (Clause c_vars' c_eq' [])
      where
        removeCondition :: ZEquality -> Solver [ZTermSubstitution]
        removeCondition (c_l, c_r) = do
          g_conds <- clauseConditions <$> sGoal <$> ask
          return $ do
            (g_l, g_r) <- g_conds
            let uni_l = unify c_l g_l
                uni_r = unify c_r g_r
                uni = uni_l `Set.union` uni_r
            if isOneToMany uni && validUnifier uni
              then [uni]
              else []
          where 
            validUnifier :: ZTermSubstitution -> Bool
            validUnifier = all (flip elem (map Var c_vars) . fst)
          
fullGoal :: Solver ZClause
fullGoal = sGoal <$> ask
            
addStep :: String -> [ZClause] -> Solver ZProof -> Solver ZProof
addStep name added solver = do
  goal <- fullGoal
  PL.addStep (PL.ProofStep name (Just goal) added) <$> solver
  
cannotDisprove :: Functor f => f [ZProof] -> f [ZProof]
cannotDisprove = fmap (filter (not . PL.disproven))
            
-- Check whether the goal is trivially true or false
checkGoal :: Solver ZProof -> Solver ZProof
checkGoal solver = do
  goal@(Clause [] (t1, t2) conds) <- sGoal <$> ask
  steps <- sConSteps <$> ask
  contr2 <- contradiction2 conds steps
  contr1 <- contradiction1 conds steps
  -- Check whether the conditions are contradictory
  if contr2
      then return $ Proven (PL.ProofStep "con" Nothing []) []
    else if PL.proven contr1
      then return $ PL.addStep (PL.ProofStep "con" Nothing []) contr1
    else if t1 == t2
      then return $ Proven (PL.ProofStep "eql" Nothing []) []
    else if unequalTerms t1 t2
      then return Unknown
      --then return $ Disproven (PL.ProofStep "irf" Nothing []) []
    else do
      -- Filter out trivial conditions
      let conds' = filter (uncurry (/=)) conds
      localGoal (Clause [] (t1, t2) conds') solver
            
contradiction1 :: [ZEquality] -> Int -> Solver ZProof
contradiction1 eqs depth 
  | depth <= 0 = return Unknown
  | otherwise = do
      inverses <- mapM solveInverse (pickEach eqs)
      return (PL.anyProven inverses)
  where 
    solveInverse :: (ZEquality, [ZEquality]) -> Solver ZProof
    solveInverse (eq@(t1, t2), others)
      | not (constructorTerm t2) = return Unknown
      | otherwise = do
        let ((Var con) : _) = flattenTerm t2
            ((TypeVar dtype) : _) = flattenAppType (termType t2)
            conterms = map Var $ filter (/= con) (zDataTypeCons dtype)
            conterms' = filter varTypeTerm conterms
        simpleSolver $ PL.anyProven <$> mapM
          (\t2' -> localGoal (Clause [] (t1, t2') others) solve) conterms'

    simpleSolver :: Solver a -> Solver a
    simpleSolver = local $ \s -> s 
      { sConSteps = 0,
        sIndSteps = depth,
        sHypotheses = mempty,
        sDoGeneralisation = True }

-- Splitting based contradition finder
--   probably subsumed by the new solver based one
contradiction2 :: [ZEquality] -> Int -> Solver Bool
contradiction2 eqs depth = do
  let (terms, kterms) = unzip eqs
  terms' <- liftZenoReader $ 
    map fst <$> mapM (applyTermDefinitions []) terms
  let eqs' = (terms' `zip` kterms)
  if any (uncurry unequalTerms) eqs' 
      then return True
    else if all (uncurry (==)) eqs' || depth <= 0
      then return False
    else do
      term_splits <- concatMapM splits (filter varTypeTerm terms')
      results <- mapM (splitOn eqs') term_splits
      return (or results)
  where
    splitOn :: [ZEquality] -> ZTerm -> Solver Bool
    splitOn eqs term
      | not (destructibleTerm term) = return False
      | otherwise =
          and <$> proveOverConstructors testSplit (termType term)
      where
        testSplit :: (ZTerm, [ZVar]) -> Solver Bool
        testSplit (term', _) = do
          let sub = substituteWithin (Set.singleton (term, term'))
          contradiction2 (map (first sub) eqs) (depth - 1)
    
    splits :: ZTerm -> Solver [ZTerm]
    splits var@(Var (ZVar _ _ ZVariable)) = return [var]
    splits term = do
      maybe_expr <- liftZenoReader $ applyFunction [] term
      case maybe_expr of
        Nothing -> return []
        Just (Value term') -> return [term']
        Just (CaseOf term' _) -> do
          terms <- splits term'
          return (term : terms)      
          
pureRecordType :: ZTypeVar -> Bool
pureRecordType (ZPolyType _) = False
pureRecordType dtype@(ZDataType _ _ cons) = 
  not (recursiveTypeVar dtype) && length cons == 1
           
invalidTerm :: ZTerm -> Bool
invalidTerm term = 
  let ((Var func) : args) = flattenTerm term
      arg_types = flattenFunType (zVarType func)
      res_types = drop (length args) arg_types 
  in null res_types
       
randomSequence :: [Solver a] -> Solver [a]
randomSequence solvers = do
  random <- sRandom <$> ask
  let (solvers', random') = shuffle solvers random
  localRandom random' (splitGens solvers')
  where
    localRandom :: StdGen -> Solver a -> Solver a
    localRandom gen = local (\s -> s { sRandom = gen })
        
    splitGens :: [Solver a] -> Solver [a]
    splitGens = foldl' splitGen (return [])
      where 
        splitGen :: Solver [a] -> Solver a -> Solver [a]
        splitGen ss s = do
          random <- sRandom <$> ask
          let (r1, r2) = split random
          ss' <- localRandom r2 ss
          s' <- localRandom r1 s
          return (s' : ss')
          
substituteTerms :: ZTermSubstitution -> Solver a -> Solver a
substituteTerms sub = 
  let f :: WithinTraversable ZTerm (t ZVar) => t ZVar -> t ZVar
      f = substituteWithin sub 
  in local $ \s -> s
      { sGoal = f (sGoal s), 
        sHypotheses = map f (sHypotheses s) }
        
enableGeneralisation :: Solver a -> Solver a
enableGeneralisation = local $ \s -> s { sDoGeneralisation = True }

disableGeneralisation :: Solver a -> Solver a
disableGeneralisation = local $ \s -> s { sDoGeneralisation = False }
    
localGoal :: ZClause -> Solver a -> Solver a
localGoal goal = local $ \s -> s { sGoal = goal }

localGoalEquality :: ZEquality -> Solver a -> Solver a
localGoalEquality eq = local $ \s -> s 
  { sGoal = (sGoal s) { clauseEquality = eq } }

localGoalConditions :: [ZEquality] -> Solver a -> Solver a
localGoalConditions conds = local $ \s -> s
  { sGoal = (sGoal s) { clauseConditions = conds } }
  
addCond :: ZEquality -> Solver a -> Solver a
addCond cond = local $ \s -> s
  { sGoal = (sGoal s) 
    { clauseConditions = cond : (clauseConditions (sGoal s)) } }

addConditions :: [ZEquality] -> Solver a -> Solver a
addConditions = foldl (\f c -> f . addCond c) id
    
addHypothesis :: ZClause -> Solver a -> Solver a
addHypothesis cls = local $ \s -> s
  { sHypotheses = cls : (sHypotheses s) }
  
removeHypothesis :: ZClause -> Solver a -> Solver a
removeHypothesis cls = local $ \s -> s
  { sHypotheses = delete cls (sHypotheses s) }
  
removeHypotheses :: [ZClause] -> Solver a -> Solver a
removeHypotheses = foldl (\f c -> f . removeHypothesis c) id

removeAllHypotheses :: Solver a -> Solver a
removeAllHypotheses = local $ \s -> s { sHypotheses = [] }

addHypotheses :: [ZClause] -> Solver a -> Solver a
addHypotheses = foldl (\f c -> f . addHypothesis c) id
  
debugShowGoal :: Solver a -> Solver a
debugShowGoal solver = do
  goal <- sGoal <$> ask
  shown <- liftZenoReader $ showZClause goal
  traceZ shown $ solver
  
generalisation :: Solver [ZProof] -> Solver [ZProof]
generalisation step = do
  do_it <- sDoGeneralisation <$> ask
  if do_it then step else return [Unknown]

takeIndStep :: Solver [ZProof] -> Solver [ZProof]
takeIndStep solver = do
  inds <- sIndSteps <$> ask
  if inds <= 0
    then return []
    else local (\s -> s { sIndSteps = inds - 1 }) solver
            
proofstep :: String -> [ZClause] -> Solver ZProofStep
proofstep name added = do
  goal <- fullGoal
  return (PL.ProofStep name (Just goal) added) 
